import archs
import torch
from torchsummary import summary
from UNet import UNet

# model = UNet(3, 4).cuda()
# a = torch.randn([1, 3, 1024, 1024]).cuda()
# model = torch.nn.Conv2d(8, 16, 3, bias=False).cuda()
# model = torch.nn.Linear(1024, 1024, bias=False).cuda()
# torchsummary.summary(torch.nn.Conv2d(8, 16, 3, bias=False).cuda(), (3, 128, 128, 128))
# summary(UNet(3, 4).cuda(), (3, 128, 128))
a = torch.nn.BatchNorm2d(385).cuda()
print(torch.cuda.memory_allocated())
